import json
from openai import OpenAI
import os
import mysql.connector
from mysql.connector import Error
import time

# 初始化
client = OpenAI(api_key="sk-427136e7d3bd47099bf14ae993d25b91", base_url="https://api.deepseek.com")
DB_CONFIG = {
    "host": "localhost",
    "port": 3306,
    "user": "root",
    "password": "",  # 替换为你的密码
    "database": "stock_db",  # 修正：应该是 'database' 而不是 'db'
    "charset": "utf8mb4"
}


def analyze_company_with_ai(stock_data):
    """
    使用DeepSeek分析上市公司数据
    """
    system_prompt = """
    你是一个金融数据分析专家。请根据提供的上市公司数据，
    提取并分析公司的基本信息、行业板块和管理者信息。
    以JSON格式输出分析结果。
    """

    user_prompt = f"""
    请分析以下上市公司数据：
    {stock_data}
    
    请提取以下信息：
    1、公司成立日期
    2、公司管理者姓名以及简介
    3、公司注册地
    4、公司性质(私有、国企、央企)
    5、一级行业板块分类
    6、二级行业板块分类
    7、公司简介经营范围
    8、公司主营业务

    请以JSON格式返回，包含以下字段：
    - establishment_date: 公司成立日期
    - manager_name: 管理者姓名
    - manager_intro: 管理者简介
    - registered_address: 注册地址
    - company_nature: 公司性质
    - primary_industry: 一级行业
    - secondary_industry: 二级行业
    - business_scope: 经营范围
    - main_business: 主营业务
    """
    print(user_prompt)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]

    try:
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            response_format={'type': 'json_object'}
        )
        # 添加return语句
        return json.loads(response.choices[0].message.content)
    except Exception as e:
        print(f"AI分析失败: {e}")
        return None


def simple_batch_analyze():
    """
    简化版的批量分析
    """
    try:
        # 数据库连接
        conn = mysql.connector.connect(**DB_CONFIG)

        # 获取股票列表
        cursor = conn.cursor(dictionary=True)
        cursor.execute("SELECT code, name FROM a_stock_basic_info")
        stocks = cursor.fetchall()

        print(f"共获取到 {len(stocks)} 只股票")


    except Error as e:
        print(f"数据库连接失败: {e}")
    except Exception as e:
        print(f"程序执行出错: {e}")
    finally:
        if 'conn' in locals() and conn.is_connected():
            cursor.close()
            conn.close()


def test_single_analysis():
    """
    测试单个股票分析
    """
    test_code = "000039"
    test_name = "测试公司"
    print(f"测试分析股票: {test_code}")

    result = analyze_company_with_ai(f"股票代码：{test_code}")
    if result:
        print("分析结果:")
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print("分析失败")


if __name__ == '__main__':
    # 可以选择运行测试或批量分析

    # test_single_analysis()  # 先测试单个

    simple_batch_analyze()  # 运行批量分析